import java.util.*;

/**
 * 398. 随机数索引
 * https://leetcode-cn.com/problems/random-pick-index/
 */
public class Solutions_398 {
    public static void main(String[] args) {
        int[] nums = {1, 2, 3, 3, 3};
        Solution_398_2 obj = new Solution_398_2(nums);
        int res1 = obj.pick(3);  // 2 or 3 or 4
        System.out.println(res1);
        int res2 = obj.pick(1);  // 0
        System.out.println(res2);
    }
}

/**
 * 解法二：蓄水池抽样（65ms）
 */
class Solution_398_2 {
    int[] nums;
    Random random = null;

    public Solution_398_2(int[] nums) {
        this.nums = nums;
        random = new Random();
    }

    public int pick(int target) {
        int cnt = 0, res = 0;
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == target) {
                cnt ++;
                // 随机生成一个 int 值，若随机数 % cnt == 0，说明被幸运地选中了
                // 关键：每个 target 都有 1/n 的机会被选取
//                if (random.nextInt() % cnt == 0) {
                if (random.nextInt(cnt) == 0) {
                    res = i;
                }
            }
        }
        return res;
    }
}

/**
 * 解法一：哈希表 + 随机数（105ms）
 */
class Solution_398_1 {
    // key：元素，value：元素索引
    Map<Integer, List<Integer>> map = null;

    public Solution_398_1(int[] nums) {
        map = new HashMap<>();
        // 添加到 map 中
        for (int i = 0; i < nums.length; i++) {
            if (!map.containsKey(nums[i])) {
                map.put(nums[i], new ArrayList<>());
            }
            map.get(nums[i]).add(i);
        }
    }

    public int pick(int target) {
        List<Integer> list = map.get(target);
        int randomIdx = (int) (Math.random() * list.size());
        return list.get(randomIdx);
    }
}

